import psycopg2
import psycopg2.extras
import sqlite3
import mysql.connector
from flask import current_app
import json

class DBService:
    @staticmethod
    def get_connection(credentials):
        """根据凭证获取数据库连接"""
        try:
            db_type = credentials.get('type', '').lower()
            
            if db_type == 'postgresql':
                conn = psycopg2.connect(
                    host=credentials.get('host'),
                    port=credentials.get('port'),
                    database=credentials.get('database'),
                    user=credentials.get('username'),
                    password=credentials.get('password')
                )
                return conn
            elif db_type == 'mysql':
                conn = mysql.connector.connect(
                    host=credentials.get('host'),
                    port=credentials.get('port'),
                    database=credentials.get('database'),
                    user=credentials.get('username'),
                    password=credentials.get('password')
                )
                return conn
            elif db_type == 'sqlite':
                conn = sqlite3.connect(credentials.get('file_path'))
                return conn
            else:
                raise ValueError(f"不支持的数据库类型: {db_type}")
        except Exception as e:
            current_app.logger.error(f"数据库连接失败: {str(e)}")
            raise
    
    @staticmethod
    def get_connection_status(credentials):
        """获取数据库连接状态"""
        try:
            conn = DBService.get_connection(credentials)
            
            if conn:
                # 尝试执行简单查询
                cursor = conn.cursor()
                
                if credentials.get('type').lower() == 'postgresql':
                    cursor.execute('SELECT version();')
                elif credentials.get('type').lower() == 'mysql':
                    cursor.execute('SELECT version();')
                else:  # SQLite
                    cursor.execute('SELECT sqlite_version();')
                    
                version = cursor.fetchone()[0]
                cursor.close()
                conn.close()
                
                return {
                    "success": True,
                    "message": "数据库连接正常",
                    "data": {
                        "status": "connected",
                        "version": version
                    }
                }
        except Exception as e:
            current_app.logger.error(f"获取数据库状态失败: {str(e)}")
            return {
                "success": False,
                "message": f"数据库连接失败: {str(e)}",
                "data": {
                    "status": "disconnected",
                    "error": str(e)
                }
            }
    
    @staticmethod
    def execute_query(credentials, query):
        """执行SQL查询"""
        try:
            conn = DBService.get_connection(credentials)
            cursor = conn.cursor()
            
            cursor.execute(query)
            
            # 获取列名
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
            
            # 获取数据
            rows = cursor.fetchall()
            
            # 将数据转换为字典列表
            data = []
            for row in rows:
                row_dict = {}
                for i, value in enumerate(row):
                    row_dict[columns[i]] = value
                data.append(row_dict)
            
            cursor.close()
            conn.close()
            
            return {
                "success": True,
                "message": "查询执行成功",
                "data": {
                    "columns": columns,
                    "rows": data
                }
            }
        except Exception as e:
            current_app.logger.error(f"执行查询失败: {str(e)}")
            return {
                "success": False,
                "message": f"查询执行失败: {str(e)}",
                "data": None
            }
    
    @staticmethod
    def execute_update(credentials, sql):
        """执行SQL更新语句"""
        try:
            conn = DBService.get_connection(credentials)
            cursor = conn.cursor()
            
            cursor.execute(sql)
            rows_affected = cursor.rowcount
            
            conn.commit()
            cursor.close()
            conn.close()
            
            return {
                "success": True,
                "message": "更新执行成功",
                "data": {
                    "rows_affected": rows_affected
                }
            }
        except Exception as e:
            current_app.logger.error(f"执行更新失败: {str(e)}")
            return {
                "success": False,
                "message": f"更新执行失败: {str(e)}",
                "data": None
            }
    
    @staticmethod
    def get_tables(credentials):
        """获取数据库表列表"""
        try:
            conn = DBService.get_connection(credentials)
            cursor = conn.cursor()
            
            db_type = credentials.get('type').lower()
            
            if db_type == 'postgresql':
                query = """
                    SELECT 
                        table_name,
                        table_type,
                        pg_catalog.obj_description(pgc.oid, 'pg_class') as table_comment
                    FROM information_schema.tables t
                    LEFT JOIN pg_catalog.pg_class pgc ON t.table_name = pgc.relname
                    WHERE table_schema = 'public'
                    ORDER BY table_name;
                """
            elif db_type == 'mysql':
                query = """
                    SELECT 
                        table_name,
                        table_type,
                        table_comment
                    FROM information_schema.tables 
                    WHERE table_schema = %s
                    ORDER BY table_name;
                """
                cursor.execute(query, (credentials.get('database'),))
            else:  # SQLite
                query = """
                    SELECT 
                        name as table_name,
                        'table' as table_type,
                        '' as table_comment
                    FROM sqlite_master
                    WHERE type='table'
                    ORDER BY name;
                """
            
            if db_type != 'mysql':  # PostgreSQL和SQLite直接执行查询
                cursor.execute(query)
            
            tables = []
            for row in cursor.fetchall():
                tables.append({
                    'name': row[0],
                    'type': row[1],
                    'comment': row[2] if row[2] else ''
                })
            
            cursor.close()
            conn.close()
            
            return {
                "success": True,
                "message": "获取表列表成功",
                "data": tables
            }
        except Exception as e:
            current_app.logger.error(f"获取表列表失败: {str(e)}")
            return {
                "success": False,
                "message": f"获取表列表失败: {str(e)}",
                "data": None
            }
    
    @staticmethod
    def get_table_structure(credentials, table_name):
        """获取表结构"""
        try:
            conn = DBService.get_connection(credentials)
            cursor = conn.cursor()
            
            db_type = credentials.get('type').lower()
            
            if db_type == 'postgresql':
                query = """
                    SELECT 
                        column_name,
                        data_type,
                        character_maximum_length,
                        column_default,
                        is_nullable,
                        col_description(pgc.oid, cols.ordinal_position) as column_comment
                    FROM information_schema.columns cols
                    LEFT JOIN pg_catalog.pg_class pgc ON cols.table_name = pgc.relname
                    WHERE table_name = %s
                    ORDER BY ordinal_position;
                """
            elif db_type == 'mysql':
                query = """
                    SELECT 
                        column_name,
                        data_type,
                        character_maximum_length,
                        column_default,
                        is_nullable,
                        column_comment
                    FROM information_schema.columns
                    WHERE table_schema = %s AND table_name = %s
                    ORDER BY ordinal_position;
                """
                cursor.execute(query, (credentials.get('database'), table_name))
            else:  # SQLite
                query = f"PRAGMA table_info('{table_name}');"
            
            if db_type == 'postgresql':
                cursor.execute(query, (table_name,))
            elif db_type == 'sqlite':
                cursor.execute(query)
            
            columns = []
            if db_type == 'sqlite':
                for row in cursor.fetchall():
                    columns.append({
                        'name': row[1],
                        'type': row[2],
                        'length': None,
                        'default': row[4],
                        'nullable': not row[3],
                        'comment': ''
                    })
            else:
                for row in cursor.fetchall():
                    columns.append({
                        'name': row[0],
                        'type': row[1],
                        'length': row[2],
                        'default': row[3],
                        'nullable': row[4] == 'YES',
                        'comment': row[5] if row[5] else ''
                    })
            
            # 获取主键信息
            if db_type == 'postgresql':
                query = """
                    SELECT a.attname
                    FROM pg_index i
                    JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                    WHERE i.indrelid = %s::regclass AND i.indisprimary;
                """
                cursor.execute(query, (table_name,))
                primary_keys = [row[0] for row in cursor.fetchall()]
            elif db_type == 'mysql':
                query = """
                    SELECT column_name
                    FROM information_schema.key_column_usage
                    WHERE table_schema = %s 
                    AND table_name = %s 
                    AND constraint_name = 'PRIMARY'
                    ORDER BY ordinal_position;
                """
                cursor.execute(query, (credentials.get('database'), table_name))
                primary_keys = [row[0] for row in cursor.fetchall()]
            else:  # SQLite
                query = f"PRAGMA table_info('{table_name}');"
                cursor.execute(query)
                primary_keys = [row[1] for row in cursor.fetchall() if row[5]]  # SQLite中pk字段为1表示主键
            
            # 获取索引信息
            if db_type == 'postgresql':
                query = """
                    SELECT
                        i.relname as index_name,
                        array_agg(a.attname) as column_names,
                        ix.indisunique as is_unique
                    FROM
                        pg_class t,
                        pg_class i,
                        pg_index ix,
                        pg_attribute a
                    WHERE
                        t.oid = ix.indrelid
                        AND i.oid = ix.indexrelid
                        AND a.attrelid = t.oid
                        AND a.attnum = ANY(ix.indkey)
                        AND t.relkind = 'r'
                        AND t.relname = %s
                    GROUP BY
                        i.relname,
                        ix.indisunique
                    ORDER BY
                        i.relname;
                """
                cursor.execute(query, (table_name,))
                indexes = []
                for row in cursor.fetchall():
                    indexes.append({
                        'name': row[0],
                        'columns': row[1],
                        'unique': row[2]
                    })
            elif db_type == 'mysql':
                query = """
                    SELECT
                        index_name,
                        GROUP_CONCAT(column_name) as column_names,
                        NOT non_unique as is_unique
                    FROM information_schema.statistics
                    WHERE table_schema = %s
                    AND table_name = %s
                    GROUP BY index_name, non_unique
                    ORDER BY index_name;
                """
                cursor.execute(query, (credentials.get('database'), table_name))
                indexes = []
                for row in cursor.fetchall():
                    indexes.append({
                        'name': row[0],
                        'columns': row[1].split(','),
                        'unique': bool(row[2])
                    })
            else:  # SQLite
                query = f"PRAGMA index_list('{table_name}');"
                cursor.execute(query)
                indexes = []
                for idx in cursor.fetchall():
                    index_name = idx[1]
                    is_unique = bool(idx[2])
                    
                    # 获取索引列
                    cursor.execute(f"PRAGMA index_info('{index_name}');")
                    columns = [row[2] for row in cursor.fetchall()]
                    
                    indexes.append({
                        'name': index_name,
                        'columns': columns,
                        'unique': is_unique
                    })
            
            cursor.close()
            conn.close()
            
            return {
                "success": True,
                "message": "获取表结构成功",
                "data": {
                    "columns": columns,
                    "primary_keys": primary_keys,
                    "indexes": indexes
                }
            }
        except Exception as e:
            current_app.logger.error(f"获取表结构失败: {str(e)}")
            return {
                "success": False,
                "message": f"获取表结构失败: {str(e)}",
                "data": None
            }
    
    @staticmethod
    def test_connection(connection_data):
        """测试数据库连接"""
        try:
            conn = DBService.get_connection(connection_data)
            
            if conn:
                cursor = conn.cursor()
                
                if connection_data.get('type').lower() == 'postgresql':
                    cursor.execute('SELECT version();')
                elif connection_data.get('type').lower() == 'mysql':
                    cursor.execute('SELECT version();')
                else:  # SQLite
                    cursor.execute('SELECT sqlite_version();')
                    
                version = cursor.fetchone()[0]
                cursor.close()
                conn.close()
                
                return {
                    "success": True,
                    "message": "数据库连接测试成功",
                    "data": {
                        "version": version
                    }
                }
        except Exception as e:
            current_app.logger.error(f"数据库连接测试失败: {str(e)}")
            return {
                "success": False,
                "message": f"数据库连接测试失败: {str(e)}",
                "data": None
            } 